import checker as chk
from collections import deque

StartSymbol = []
ProdRule = {}
ProdType = {}
BfsQueue = deque()
MRQueue = deque()
FuncDefineStr = None
Ck = None

ArrayIndexOffset = []
ArrayElementArgs = []
ArrayAllowed = []
ReduceTaskLen = 0
ARRAY_IDX = "!!IDX"
ARRAY_PREV = "!!PRV"
ARRAY_CUR = "!!CUR"
ARRAY_NO_INITIAL = "!!NO_INIT"


def get_min_nt_depth(stmt):
    if type(stmt) == list:
        depth = 0
        for sub_stmt in stmt:
            xd = get_min_nt_depth(sub_stmt)
            if xd > 0 and (depth == 0 or xd + 1 < depth):
                depth = xd + 1
        return depth
    else:
        return 1 if stmt in ProdRule else 0


def extend_text_at_depth(stmt, depth, r_type, is_r_top):
    result = []
    if depth == 0:
        return []
    if is_r_top:
        try_extend = extend_text_at_depth(stmt[0], depth - 1, [ARRAY_CUR, ARRAY_PREV, ARRAY_IDX], False)
        if len(try_extend) > 0:
            for extended in try_extend:
                result.append([extended, stmt[1]])
        else:
            if depth > 2:
                depth = 2
            try_extend = extend_text_at_depth(stmt[1], depth - 1, [], False)
            for extended in try_extend:
                result.append([stmt[0], extended])
        return result
    if type(stmt) == list:
        for i in reversed(range(len(stmt))):
            try_extend = extend_text_at_depth(stmt[i], depth - 1, r_type, False)
            if len(try_extend) > 0:
                for extended in try_extend:
                    result.append(stmt[0:i] + [extended] + stmt[i + 1:])
                break
    elif stmt in ProdRule:
        if r_type is None:
            for extended in ProdRule[stmt]:
                result.append(extended)
        else:
            for extended in ProdRule[stmt]:
                if is_non_empty_list(extended) and extended[0] == 'ite' and ARRAY_CUR not in r_type:
                    continue
                if extended not in ArrayElementArgs:
                    result.append(extended)
            if ProdType[stmt] == 'Int':
                for extended in r_type:
                    if extended in ArrayAllowed:
                        result.append(extended)
    return result


###########################
# helpers
###########################

def is_non_empty_list(stmt, min_len=1):
    return type(stmt) == list and len(stmt) >= min_len


def operand_is_str(*ss):
    for s in ss:
        if type(s) != str or s in ProdRule:
            return False
    return True


def string_is_int(s):
    return s != '+' and s != '-' and s.lstrip('+-').isdigit()


def stmt_is_ite(stmt):
    return is_non_empty_list(stmt) and stmt[0] == 'ite'


###################################################
# prune rules
# WARNING:
#    PRECONDITION: stmt is list with proper length
###################################################

def forbid_double_not(stmt):
    return stmt[0] in ['not', 'ite'] and \
           is_non_empty_list(stmt[1]) and \
           stmt[1][0] == 'not'


def forbid_sym_bin_op_1(stmt):
    return stmt[0] in ['+', '=', '*'] and \
           is_non_empty_list(stmt[1]) and \
           operand_is_str(stmt[2])


def forbid_sym_bin_op_2(stmt):
    return stmt[0] in ['+', '=', '*'] and \
           operand_is_str(stmt[1], stmt[2]) and \
           stmt[1] > stmt[2]


def forbid_bin_op_nested_ite(stmt):
    return stmt[0] in ['+', '-', '*', '/', 'mod'] and \
           (stmt_is_ite(stmt[1]) or stmt_is_ite(stmt[2]))


def forbid_sym_compare(stmt):
    return stmt[0] in ['<', '>', '<=', '>=', '='] and \
           operand_is_str(stmt[1], stmt[2]) and \
           stmt[1] == stmt[2]


def forbid_trivial_compare(stmt):
    return stmt[0] in ['<', '>', '<=', '>=', '='] and \
           operand_is_str(stmt[1], stmt[2]) and \
           string_is_int(stmt[1]) and \
           string_is_int(stmt[2])


def forbid_noop_plus(stmt):
    return stmt[0] == '+' and (stmt[1] == '0' or stmt[2] == '0')


def forbid_noop_minus(stmt):
    return stmt[0] == '-' and stmt[2] == '0'


####################################
# Check whether prune
####################################


def forbidden_sub_expr(stmt):
    if type(stmt) == list:
        for sub_stmt in stmt:
            if forbidden_expr(sub_stmt):
                return True
    return False


def forbidden_curr_expr(stmt):
    if not is_non_empty_list(stmt, 2):
        return False
    if forbid_double_not(stmt):
        return True
    return is_non_empty_list(stmt, 3) and (
            forbid_bin_op_nested_ite(stmt) or
            forbid_sym_bin_op_1(stmt) or
            forbid_sym_bin_op_2(stmt) or
            forbid_sym_compare(stmt) or
            forbid_trivial_compare(stmt) or
            forbid_noop_plus(stmt) or
            forbid_noop_minus(stmt))


def forbidden_expr(stmt):
    return forbidden_curr_expr(stmt) or forbidden_sub_expr(stmt)


################################
# operator deduplication
################################

def lst_get_op(lst, op):
    for i in range(len(lst)):
        stmt = lst[i]
        if type(stmt) == list and len(stmt) == 3 and stmt[0] == op and stmt[1] in ProdRule and stmt[2] in ProdRule:
            return i
    return None


def lst_rm_op(lst, op):
    while True:
        i = lst_get_op(lst, op)
        if i is None:
            return
        del lst[i]


#######################################
# convert to simple list
#######################################

def flatten_type_annotated_list(lst):
    if type(lst) == tuple:
        return str(lst[1])
    if type(lst) == list:
        return list([flatten_type_annotated_list(sub) for sub in lst])
    return str(lst)


######################################
# map reduce task
######################################
def func_arg_sanity_check(arg_def, start_type):
    if start_type != 'Int' or len(StartSymbol) != 1:
        return False
    for a in arg_def:
        if a[1] != 'Int' or a[0] not in ProdRule[StartSymbol[0]]:
            return False
    return True


def is_indexed_search_task(a_args):
    if len(a_args) < 3:
        return 0
    arg_prefix = a_args[0][:-1]
    start_idx = 0 if a_args[0][-1] == '0' else 1
    for i in range(len(a_args)):
        expected_name = arg_prefix + str(i + start_idx)
        if expected_name != a_args[i]:
            return 0
    task_len = len(a_args)
    ok_idx = [True, True]
    for i in range(len(a_args)):
        for offset in [0, 1]:
            if str(i + offset) not in ProdRule[StartSymbol[0]]:
                ok_idx[offset] = False
    for offset in [0, 1]:
        if ok_idx[offset]:
            ArrayIndexOffset.append(offset)
    if len(ArrayIndexOffset) == 0:
        return 0
    for a in a_args:
        ArrayElementArgs.append(a)
    return task_len


def is_ordinary_reduce_task(a_args):
    if len(a_args) < 3:
        return 0
    for a in a_args:
        ArrayElementArgs.append(a)
    return len(a_args)


def expand_reduce_stmt(stmt, cur, idx, prv):
    if type(stmt) == list:
        return list([expand_reduce_stmt(sub, cur, idx, prv) for sub in stmt])
    elif stmt == ARRAY_CUR:
        return cur
    elif stmt == ARRAY_IDX:
        return str(idx)
    elif stmt == ARRAY_PREV:
        return prv
    return stmt


######################################
# main
######################################

def fallback_prepare(bm_expr):
    global FuncDefineStr
    global ReduceTaskLen
    global Ck
    Ck = chk.Checker(bm_expr)
    syn_fun_expr = []
    for expr in bm_expr:
        if len(expr) == 0:
            continue
        elif expr[0] == 'synth-fun':
            syn_fun_expr = expr
    func_def = ['define-fun'] + syn_fun_expr[1:4]
    FuncDefineStr = chk.list_to_string(func_def)
    start_type = syn_fun_expr[3]
    fun_arg_def = syn_fun_expr[2]
    fun_arg_name = list([a[0] for a in fun_arg_def])

    for nt_full_sig in syn_fun_expr[4]:
        nt_name = nt_full_sig[0]
        nt_type = nt_full_sig[1]
        if nt_type == start_type:
            StartSymbol.append(nt_name)
            BfsQueue.append((nt_name, 1))
        ProdType[nt_name] = nt_type
        ProdRule[nt_name] = []

    for nt_full_sig in syn_fun_expr[4]:
        nt_name = nt_full_sig[0]
        for nt in nt_full_sig[2]:
            ProdRule[nt_name].append(flatten_type_annotated_list(nt))

    ReduceTaskLen = 0
    if len(fun_arg_name) > 3 and func_arg_sanity_check(fun_arg_def, start_type):
        ReduceTaskLen = is_indexed_search_task(fun_arg_name)
        if ReduceTaskLen == 0:
            ReduceTaskLen = is_indexed_search_task(fun_arg_name[:-1])
        if ReduceTaskLen == 0:
            ReduceTaskLen = is_indexed_search_task(fun_arg_name[1:])
        if ReduceTaskLen == 0:
            ReduceTaskLen = is_ordinary_reduce_task(fun_arg_name)
    if ReduceTaskLen > 0:
        start_c1 = [StartSymbol[0], ARRAY_NO_INITIAL]
        start_c2 = [StartSymbol[0], StartSymbol[0]]
        MRQueue.append((start_c1, get_min_nt_depth(start_c1)))
        MRQueue.append((start_c2, get_min_nt_depth(start_c2)))
        ArrayAllowed.append(ARRAY_CUR)
        ArrayAllowed.append(ARRAY_PREV)
        if len(ArrayIndexOffset) == 0:
            ArrayIndexOffset.append(-99999)  # placeholder
        else:
            ArrayAllowed.append(ARRAY_IDX)
    # print(MRQueue[0])
    # print(MapReduceTaskLen)
    # print(ArrayElementArgs)
    # print(ArrayIndexOffset)
    # print(ArrayAllowed)
    # raise NotImplementedError()


def fallback_iteration():
    if len(BfsQueue) == 0:
        return None
    curr, cd = BfsQueue.popleft()
    if forbidden_expr(curr):
        return None
    if cd == 0:
        curr_str = chk.list_to_string(curr)
        full_str = FuncDefineStr[:-1] + ' ' + curr_str + FuncDefineStr[-1]
        counterexample = Ck.get_counterexample(full_str)
        if counterexample is None:  # No counter-example
            return full_str
    else:
        try_extend = extend_text_at_depth(curr, cd, None, False)
        for te in try_extend:
            nd = get_min_nt_depth(te)
            if nd >= cd:  # DON'T CHANGE >= to > !!!
                BfsQueue.append((te, nd))
            else:
                BfsQueue.appendleft((te, nd))
        return None


def reduce_task_iteration():
    if len(MRQueue) == 0:
        return None
    curr, cd = MRQueue.popleft()
    if forbidden_expr(curr[0]):
        return None
    if cd == 0:
        for idx_offset in ArrayIndexOffset:
            last_s = ArrayElementArgs[0]
            start_idx = 1
            if curr[1] != ARRAY_NO_INITIAL:
                last_s = chk.list_to_string(curr[1])
                start_idx = 0
            for i in range(start_idx, ReduceTaskLen):
                last_s = chk.list_to_string(expand_reduce_stmt(curr[0], ArrayElementArgs[i], i + idx_offset, last_s))

            full_str = FuncDefineStr[:-1] + ' ' + last_s + FuncDefineStr[-1]
            counterexample = Ck.get_counterexample(full_str)
            if counterexample is None:  # No counter-example
                return full_str
    else:
        try_extend = extend_text_at_depth(curr, cd, [], True)
        for te in try_extend:
            nd = get_min_nt_depth(te)
            if nd >= cd:  # DON'T CHANGE >= to > !!!
                MRQueue.append((te, nd))
            else:
                MRQueue.appendleft((te, nd))
        return None
